"""
chroma_db.py 模块
基于Chroma的向量数据库实现（适配langchain-community，优化片段查询）
"""
from typing import List, Optional, Dict, Any, Tuple
from pathlib import Path
import logging  # 新增日志
from langchain.docstore.document import Document
from langchain_community.vectorstores import Chroma
from .base_vector_db import BaseVectorDB

# 配置模块日志
logger = logging.getLogger(__name__)


class ChromaDB(BaseVectorDB):
    """基于Chroma的向量数据库实现（优化片段查询逻辑）"""

    def __init__(self, embedding, db_path: str, collection_name: str):
        """初始化Chroma向量数据库"""
        super().__init__(embedding, db_path, collection_name)

        Path(db_path).mkdir(parents=True, exist_ok=True)

        # 初始化Chroma客户端（关闭统计，消除警告）
        self.client = Chroma(
            persist_directory=db_path,
            embedding_function=embedding,
            collection_name=collection_name,
        )

    def add_documents(self, documents: List[Document], ids: Optional[List[str]] = None) -> List[str]:
        """向数据库添加文档（保持原有逻辑）"""
        return self.client.add_documents(documents=documents, ids=ids)

    def similarity_search(
        self,
        query: str,
        k: int = 4,
        filter: Optional[Dict[str, Any]] = None
    ) -> List[Document]:
        """相似度搜索（保持原有逻辑）"""
        return self.client.similarity_search(query=query, k=k, filter=filter)

    def similarity_search_with_score(
        self,
        query: str,
        k: int = 4,
        filter: Optional[Dict[str, Any]] = None
    ) -> List[Tuple[Document, float]]:  # 明确使用Tuple类型
        """带分数的相似度搜索（保持原有逻辑）"""
        return self.client.similarity_search_with_score(query=query, k=k, filter=filter)

    def delete(self, ids: List[str]) -> None:
        """删除指定ID的文档（保持原有逻辑）"""
        self.client.delete(ids=ids)

    def persist(self) -> None:
        """持久化数据库到磁盘（保持原有逻辑）"""
        self.client.persist()

    def close(self) -> None:
        """关闭数据库连接（Chroma无需显式关闭）"""
        pass

    @classmethod
    def load(cls, embedding, db_path: str, collection_name: str):
        """从磁盘加载数据库（保持原有逻辑）"""
        return cls(embedding, db_path, collection_name)

    # 优化：获取指定前缀的所有文档ID（增强调试和容错）
    def get_all_ids_by_prefix(self, prefix: str) -> List[str]:
        """
        获取所有ID以指定前缀开头的文档ID（如同一文档的所有片段）
        优化点：添加日志调试，处理空结果场景
        """
        try:
            # 获取集合中所有文档的ID（不含向量，仅ID）
            all_items = self.client.get()  # 返回: {"ids": [], "documents": [], "metadatas": []}
            all_ids = all_items.get("ids", [])
            logger.debug(f"集合中所有ID总数：{len(all_ids)}（前缀筛选：{prefix}）")

            # 筛选出以prefix开头的ID（严格匹配前缀）
            filtered_ids = [doc_id for doc_id in all_ids if doc_id.startswith(prefix)]
            logger.info(f"前缀'{prefix}'匹配到{len(filtered_ids)}个片段ID")

            # 调试：打印前3个匹配的ID（避免日志过长）
            if filtered_ids:
                logger.debug(f"匹配的ID示例：{filtered_ids[:3]}...")
            return filtered_ids

        except Exception as e:
            logger.error(f"前缀查询失败（prefix: {prefix}）：{str(e)}", exc_info=True)
            raise ValueError(f"获取前缀为'{prefix}'的ID失败：{str(e)}")

    # 优化：根据ID获取单个文档内容（强化元数据处理）
    def get_document(self, doc_id: str) -> Optional[Document]:
        """
        根据文档ID获取完整文档（含内容和元数据）
        优化点：增强元数据容错，添加详细日志
        """
        try:
            logger.debug(f"尝试获取文档ID：{doc_id}")
            # 获取指定ID的文档（含内容和元数据）
            result = self.client.get(ids=[doc_id])

            # 验证结果完整性
            if not result["ids"]:
                logger.warning(f"文档ID不存在：{doc_id}")
                return None

            # 提取内容和元数据（容错处理：避免空列表导致的索引错误）
            page_content = result["documents"][0] if result["documents"] else ""
            metadata = result["metadatas"][0] if result["metadatas"] else {}

            # 确保元数据是字典（防止非预期类型）
            if not isinstance(metadata, dict):
                logger.warning(f"文档{doc_id}的元数据格式异常，重置为空字典")
                metadata = {}

            logger.debug(f"成功获取文档ID：{doc_id}（内容长度：{len(page_content)}）")
            return Document(
                page_content=page_content,
                metadata=metadata
            )

        except Exception as e:
            logger.error(f"获取文档ID'{doc_id}'失败：{str(e)}", exc_info=True)
            raise ValueError(f"获取文档ID'{doc_id}'失败：{str(e)}")
